# -*- coding: utf-8 -*-

#    OriginDB
#    Copyright (C) 2010  Jason Power <power.jg@gmail.com>
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

from PyQt4.QtGui import *
from PyQt4 import QtCore

import openbabel
import pybel

from math import sqrt, atan2, pi, cos, sin

from atom import Atom
from bond import Bond
from minimize import GenCoords


class Compound(QGraphicsItem):

	scale = 50.0
	
	selectedPen = QPen(QColor(26,66,186))
	selectedPen.setWidth(2)
	selectedPen.setCapStyle(QtCore.Qt.FlatCap)

	selectedPadding = 7

	def __init__(self, mol=None, showHydrogens=False, parent=None):
		QGraphicsItem.__init__(self, parent)

		if mol and type(mol)==pybel.Molecule:
			mol = mol.OBMol
		elif mol and type(mol)==str:
			pymol = pybel.readstring("sdf", mol)
			mol = pymol.OBMol
		if mol:
			self.mol = mol
			gen = GenCoords()
			gen.calcCoords(mol)
			gen.setCoords()
		else:
			self.mol = openbabel.OBMol()

		self.mol.Center()

		self.atoms = []
		self.bonds = []

		self.totRot = 0
		self.thisRot = 0

		self.showHydrogens = showHydrogens
		if showHydrogens:
			self.mol.AddHydrogens()
		elif not showHydrogens:
			self.mol.DeleteHydrogens()

		self.selected = False

		self.setFlag(QGraphicsItem.ItemIsMovable, True)
		self.setFlag(QGraphicsItem.ItemIsSelectable, True)
		try:
			self.setFlag(QGraphicsItem.ItemSendsGeometryChanges, True)
		except AttributeError:
			pass #Need this in Qt4.6 not in Qt4.5

	def addAtom(self, atom, pos):
		pos = self.mapFromScene(pos)
		
		self.mol.BeginModify()
		newAtom = self.mol.NewAtom()
		newAtom.Duplicate(atom)
		newAtom.SetVector(self.translateToAtomCoords(pos))
		self.mol.EndModify()
		atomItem = Atom(newAtom, parent=self)
		atomItem.rotate(-self.totRot)
		self.atoms.append(atomItem)

		self.prepareGeometryChange()
		return newAtom

	def removeAtom(self, atom):
		print "Trying to remove", atom.GetIdx()
		idx = atom.GetIdx()
		for i in range(len(self.atoms)):
			if self.atoms[i].atom.GetIdx() == idx:
				self.scene().removeItem(self.atoms[i])
				del self.atoms[i]
				break

		pos = self.translateFromAtomCoords([atom.GetX(), atom.GetY()])

		self.mol.BeginModify()

		foundBond = True
		while foundBond:
			foundBond = False
			for i in range(len(self.bonds)):
				if self.bonds[i].bond.GetBeginAtomIdx() == atom.GetIdx() or self.bonds[i].bond.GetEndAtomIdx() == atom.GetIdx():
					self.scene().removeItem(self.bonds[i])
					print "Removing bond:", self.bonds[i].bond.GetIdx()
					self.mol.DeleteBond(self.bonds[i].bond)
					del self.bonds[i]
					foundBond = True
					break

		self.mol.DeleteAtom(atom)
		self.mol.EndModify()

		self.prepareGeometryChange()
		return pos

	def addBond(self, beginAtom, endAtom, typ):
		if typ=="single" or typ=="wedge" or typ=="hash" or typ=="squiggle":
			bo = 1
		elif typ=="double" or typ=="cistrans":
			bo = 2
		elif typ=="triple":
			bo = 3

		self.mol.BeginModify()
		self.mol.AddBond(beginAtom.GetIdx(), endAtom.GetIdx(), bo)

		bond = self.mol.GetBond(beginAtom, endAtom)
		
		if typ=="wedge":
			bond.SetWedge()
		elif typ=="hash":
			bond.SetHash()
		
		self.mol.EndModify()

		b = Bond(bond, parent=self)
		self.bonds.append(b)

		for atom in self.atoms:
			atom.updateSymbol()

		self.prepareGeometryChange()
		return bond

	def addMolecule(self, mol, pos=QtCore.QPointF(0,0)):
		pos = self.mapFromScene(pos)
		
		for i in self.atoms:
			self.scene().removeItem(i)
		for i in self.bonds:
			self.scene().removeItem(i)
		self.atoms = []
		self.bonds = []

		gen = GenCoords()
		gen.calcCoords(mol)
		gen.setCoords()
		mol.Center()

		for atom in openbabel.OBMolAtomIter(self.mol):
			atom.SetVector(atom.GetX()-pos.x()/self.scale, atom.GetY()-pos.y()/self.scale, 0)

		self.mol += mol

		self.mol.DeleteHydrogens()

		self.addToScene()
		return [], []

	def addRing(self, degree, bond=None, pos=QtCore.QPointF(0,0), double=False):
		if not bond:
			mol = openbabel.OBMol()
			atoms = []
			print degree
			for i in range(1, degree+1):
				a = mol.NewAtom()
				a.SetAtomicNum(6)
				atoms.append(a)

			for i in range(1, degree+1):
				if double:
					order = 2-i%2 # bo=2 if i%2 1 otherwise
				else:
					order = 1
				
				if i != degree:
					mol.AddBond(i, i+1, order)
				else:
					mol.AddBond(i, 1, order)

			self.addMolecule(mol, pos)

		else:
			l = bond.GetLength()

			beginAtom = bond.GetBeginAtom()
			endAtom = bond.GetEndAtom()

			x1 = beginAtom.GetX()
			y1 = beginAtom.GetY()
			x2 = endAtom.GetX()
			y2 = endAtom.GetY()

			#if x1 > x2:
				#print "flip x"
				#x1, x2 = x2, x1
				#y1, y2 = y2, y1
				#beginAtom, endAtom = endAtom, beginAtom
			#if y1 > y2:
				#print "flip y"
				#x1, x2 = x2, x1
				#y1, y2 = y2, y1
				#beginAtom, endAtom = endAtom, beginAtom

			print "begin", x1, y1
			print "end", x2, y2
			#print l

			if double:
				if bond.GetBO() == 2:
					addDouble = 1
				else:
					addDouble = 0

			interiorAngle = -pi*(degree-2)/degree
			angle = atan2(y2-y1,x2-x1)
			print angle*180/pi
			
			x, y = x2, y2
			print x1+l*cos(angle), y1+l*sin(angle)
			last = endAtom
			if double:
				doublesAdded = addDouble
			for i in range(degree-2):
				angle = interiorAngle-(pi-angle)
				print -angle*180/pi
				
				x += l*cos(-angle)
				y += l*sin(-angle)
				
				print "here:", x, y

				a = self.mol.NewAtom()
				a.SetAtomicNum(6)
				a.SetVector(x, y, 0)
				
				atomItem = Atom(a, parent=self)
				atomItem.rotate(-self.totRot)
				self.atoms.append(atomItem)
				if double:
					if not double or not i%2 == addDouble:
						self.addBond(last, a, "single")
					else:
						self.addBond(last, a, "double")
						doublesAdded += 1
				else:
					self.addBond(last, a, "single")

				last = a

			if double:
				if (not double) or (not (degree-2)%2 == addDouble) or (float(degree)/(doublesAdded+1)) < 2:
					self.addBond(beginAtom, last, "single")
				else:
					self.addBond(beginAtom, last, "double")
			else:
				self.addBond(beginAtom, last, "single")

			self.prepareGeometryChange()

			print "\n"

		
	def removeBond(self, bond):
		print "Trying to remove", bond.GetIdx()
		idx = bond.GetIdx()
		for i in range(len(self.bonds)):
			if self.bonds[i].bond.GetIdx() == idx:
				self.scene().removeItem(self.bonds[i])
				del self.bonds[i]
				break

		beginAtom = bond.GetBeginAtom()
		endAtom = bond.GetEndAtom()

		order = bond.GetBO()
		typ = ""
		if order==1:
			if bond.IsHash():
				typ = "hash"
			elif bond.IsWedge():
				typ = "wedge"
			else:
				typ = "single"
		elif order==2:
			typ = "double"
		elif order==3:
			typ = "triple"

		self.mol.BeginModify()
		self.mol.DeleteBond(bond)
		self.mol.EndModify()

		self.prepareGeometryChange()
		print "New len of bonds:", len(self.bonds)
		return beginAtom, endAtom, typ

	def setShowHydrogens(self, show):
		self.showHyrdogens = show
		if show:
			self.mol.AddHydrogens()
		elif not show:
			self.mol.DeleteHydrogens()

		for i in self.atoms:
			self.scene().removeItem(i)
		for i in self.bonds:
			self.scene().removeItem(i)
		self.atoms = []
		self.bonds = []
		self.addToScene()
		self.prepareGeometryChange()
		self.update()

	def flipAtBond(self, bond):
		if bond.IsInRing():
			return

		before = []
		after = []
		for b in openbabel.OBAtomBondIter(bond.GetBeginAtom()):
			if b.GetIdx() != bond.GetIdx() and b.GetIdx() not in before:
				before.append(b.GetIdx())
				before += self._bondDFS(b, set([bond.GetIdx()]))
		
		after = self._bondDFS(bond, set(before))

		before = set(before)
		after = set(after)
		after.remove(bond.GetIdx())

		#print len(before), [b for b in before]
		#print len(after), [b for b in after]

		if len(before) < len(after):
			working = before
		else:
			working = after

		batom = bond.GetBeginAtom()
		eatom = bond.GetEndAtom()
		# vector perp to line of bond
		v = (eatom.GetY() - batom.GetY(),
			 -(eatom.GetX() - batom.GetX()))
		vd = sqrt(v[0]*v[0]+v[1]*v[1])
		vnorm = (v[0]/vd, v[1]/vd)

		atomsDone = []
		for b in openbabel.OBMolBondIter(self.mol):
			if not b.GetIdx() in working:
				continue

			for a in [b.GetBeginAtom(), b.GetEndAtom()]:
				if a.GetIdx() in atomsDone:
					continue
				# vector from atom we are looking at to first point
				r = (batom.GetX() - a.GetX(), batom.GetY() - a.GetY())
				# dist from point to line: (vnorm dot r)
				d = vnorm[0]*r[0]+vnorm[1]*r[1]
				newPos = (a.GetX()+2*vnorm[0]*d, a.GetY()+2*vnorm[1]*d)
				a.SetVector(newPos[0], newPos[1], 0)

				atomsDone.append(a.GetIdx())

		for atom in self.atoms:
			atom.updatePos()

	def _bondDFS(self, bond, visited=set()):

		ans = []
		for atom in [bond.GetBeginAtom(), bond.GetEndAtom()]:
			itr = openbabel.OBAtomBondIter(atom)
			for b in itr:
				if b.GetIdx() == bond.GetIdx():
					continue
				if b.GetIdx() in visited or b.GetIdx() in ans:
					continue
				visited.add(b.GetIdx())

				ans += self._bondDFS(b, visited)+[b.GetIdx()]
			
		return ans

	def translateFromAtomCoords(self, coords):
		""" @coords is a 2 or 3 tuple
			Returns a QPointF
		"""
		return QtCore.QPointF(coords[0]*self.scale, coords[1]*self.scale)

	def translateToAtomCoords(self, point):
		""" @point is a QPointF
			Returns a openbabel vector3
		"""
		return pybel.ob.vector3(point.x()/self.scale, point.y()/self.scale, 0)

	def minimize(self):

		#for atom in openbabel.OBMolAtomIter(self.mol):
			#print "Atom", atom.GetIdx(), atom.GetAtomicNum(), ":", atom.GetX(), atom.GetY(), atom.GetZ()

		g = GenCoords()
		g.calcCoords(self.mol)
		g.setCoords()

		#for atom in openbabel.OBMolAtomIter(self.mol):
			#print "Atom", atom.GetIdx(), atom.GetAtomicNum(), ":", atom.GetX(), atom.GetY(), atom.GetZ()

		for atom in self.atoms:
			atom.updatePos()

	def getAtom(self, idx):
		if idx > self.mol.NumAtoms():
			raise ValueError("Idx is too large, bug!")
		raise Exception
		return

	def addToScene(self):
		for atom in openbabel.OBMolAtomIter(self.mol):
			a = Atom(atom, parent=self)
			self.atoms.append(a)

		for bond in openbabel.OBMolBondIter(self.mol):
			b = Bond(bond, parent=self)
			self.bonds.append(b)

	def updateBonds(self, atom):
		for bond in self.bonds:
			if bond.bond.GetBeginAtomIdx()==atom.GetIdx() or bond.bond.GetEndAtomIdx()==atom.GetIdx():
				bond.atomMoved()

		self.prepareGeometryChange()

	def rotate(self, deg):
		QGraphicsItem.rotate(self, deg)
		self.totRot += deg
		self.thisRot += deg

	def doneRotate(self):
		print self.thisRot
		for atom in self.atoms:
			atom.rotate(-self.thisRot)
		self.thisRot = 0

	def itemChange(self, change, value):
		if change == QGraphicsItem.ItemSceneHasChanged and value:
			self.addToScene()
		return QGraphicsItem.itemChange(self, change, value)

	def setSelected(self, selected):
		self.selected = selected
		self.update()

	def boundingRect(self):
		rect = self.childrenBoundingRect()
		rect = QtCore.QRectF(rect.x()-self.selectedPadding/2, rect.y()-self.selectedPadding/2, rect.width()+self.selectedPadding, rect.height()+self.selectedPadding)
		return rect

	def paint(self, painter, option, widget):
		if self.selected:
			painter.setPen(self.selectedPen)
			painter.drawRect(self.boundingRect())
